{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Waiting for Execution of Another Electron\n",
    "\n",
    "Covalent parallelizes execution of tasks (electrons) when their inputs and outputs are independent. You might nonetheless need an electron to wait for another electron to finish execution before starting.\n",
    "\n",
    "The Covalent `Electron` class provides a method, `wait_for()`, that explicitly defers execution until another electron has finished executing. In order to use that one can call `ct.wait()` as described below.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "Define a lattice where an electron `task_A` needs an electron `task_B` to complete before starting, even though the input of `task_A` does not depend directly on the output of `task_B`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import covalent as ct\n",
    "\n",
    "@ct.electron\n",
    "def task_A(**args_A):\n",
    "    result_A = 1\n",
    "    # Some task computation here\n",
    "    return result_A\n",
    "\n",
    "@ct.electron\n",
    "def task_B(**args_B):\n",
    "    result_B = 2\n",
    "    # Some task computation here\n",
    "    return result_B\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Procedure\n",
    "\n",
    "Use the `ct.wait()` function to force `task_A` to delay start of execution until `task_B` is completed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.lattice\n",
    "def workflow():\n",
    "    arg_A = 0\n",
    "    arg_B = 0\n",
    "    result_B = task_B(arg_B)\n",
    "    result_A = task_A(arg_A)\n",
    "    ct.wait(result_A, result_B) # Wait for result_B before computing result_A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, the call to `ct.wait(result_A, result_B)` creates the \"artificial\" dependency of `task_A` on `task_B`. The example produces this transport graph:\n",
    "\n",
    "![Transport graph: waiting edge](./images/waiting_edge.png)\n",
    "\n",
    "Note the edge between the two electrons is labeled `!waiting_edge` rather than with a parameter dependency.\n",
    "\n",
    "Note also that the `ct.wait()` call appears in the workflow after the calls to both electrons. Recall that the lattice is run (and the transport graph built) before any electrons are executed; this enables the server to instantiate the dependency before dispatching the electrons."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following example, `task_3` waits for completion of `task_1` even though `task_3`'s execution does not directly depend on the output of `task_1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: 1500\n",
      "input args: []\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-02-17 00:48:57.309061\n",
      "end_time: 2023-02-17 00:49:00.525336\n",
      "\n",
      "results_dir: /Users/mini-me/agnostiq/covalent/doc/source/how_to/coding/results\n",
      "dispatch_id: a28b546f-7493-4eff-8dcb-c6425b87c804\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): 4\n",
      ":parameter:2(1): 2\n",
      "task_2(2): 12\n",
      ":parameter:3(3): 3\n",
      "task_3(4): 125\n",
      ":parameter:5(5): 5\n",
      "task_2(6): 1500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "\n",
    "@ct.electron\n",
    "def task_1(a):\n",
    "    import time\n",
    "    time.sleep(3)\n",
    "    return a ** 2\n",
    "\n",
    "@ct.electron\n",
    "def task_2(x, y):\n",
    "    return x * y\n",
    "\n",
    "@ct.electron\n",
    "def task_3(b):\n",
    "    return b ** 3\n",
    "\n",
    "@ct.lattice\n",
    "def workflow():\n",
    "    res_1 = task_1(2)\n",
    "    res_2 = task_2(res_1, 3)\n",
    "    res_3 = task_3(5)\n",
    "    ct.wait(res_3, res_1)\n",
    "\n",
    "    return task_2(res_2, res_3)\n",
    "\n",
    "dispatch_id = ct.dispatch(workflow)()\n",
    "result = ct.get_result(dispatch_id, wait=True)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `wait()` function can take a list of electrons rather than a single electron. The calling task waits until all electrons in the list are finished before executing.\n",
    "\n",
    "In the following example, `task_3` waits for electrons 1a, 1b, and 1c before executing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: 1500\n",
      "input args: []\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-02-17 00:49:01.456059\n",
      "end_time: 2023-02-17 00:49:02.663492\n",
      "\n",
      "results_dir: /Users/mini-me/agnostiq/covalent/doc/source/how_to/coding/results\n",
      "dispatch_id: 83885751-d800-4e36-b785-efa01e58c3be\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1a(0): 4\n",
      ":parameter:2(1): 2\n",
      "task_1b(2): 8\n",
      ":parameter:2(3): 2\n",
      "task_1c(4): 16\n",
      ":parameter:2(5): 2\n",
      "task_2(6): 12\n",
      ":parameter:3(7): 3\n",
      "task_3(8): 125\n",
      ":parameter:5(9): 5\n",
      "task_2(10): 1500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "import time\n",
    "\n",
    "\n",
    "@ct.electron\n",
    "def task_1a(a):\n",
    "    time.sleep(1)\n",
    "    return a ** 2\n",
    "\n",
    "@ct.electron\n",
    "def task_1b(a):\n",
    "    time.sleep(1)\n",
    "    return a ** 3\n",
    "\n",
    "@ct.electron\n",
    "def task_1c(a):\n",
    "    time.sleep(1)\n",
    "    return a ** 4\n",
    "\n",
    "@ct.electron\n",
    "def task_2(x, y):\n",
    "    return x * y\n",
    "\n",
    "@ct.electron\n",
    "def task_3(b):\n",
    "    return b ** 3\n",
    "\n",
    "@ct.lattice\n",
    "def workflow():\n",
    "    res_1a = task_1a(2)\n",
    "    res_1b = task_1b(2)\n",
    "    res_1c = task_1c(2)\n",
    "    res_2 = task_2(res_1a, 3)\n",
    "    res_3 = task_3(5)\n",
    "    ct.wait(res_3, [res_1a, res_1b, res_1c])\n",
    "\n",
    "    return task_2(res_2, res_3)\n",
    "\n",
    "dispatch_id = ct.dispatch(workflow)()\n",
    "result = ct.get_result(dispatch_id, wait=True)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tip: Run the examples above and use the Covalent GUI to view the transport graphs.\n",
    "\n",
    "The `wait()` function is especially useful when working with a dependent task that does not take any inputs and/or when you need to wait for classes that do not return an output.\n",
    "\n",
    "### See Also\n",
    "\n",
    "[Adding an Electron to a Lattice](./add_electron_to_lattice.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
